# -*- coding: utf-8 -*-
import copy
import datetime
import logging
import time
from traceback import format_exc
from pymongo import UpdateOne, MongoClient
from scrapy.utils.project import get_project_settings
from zc_core.client.mongo_client import Mongo
from zc_core.util.batch_gen import time_to_batch_no, batch_to_year
from zc_core.util.encrypt_util import md5


class SimpleDao(object):
    db_name_tpl = 'ctaxccgp_{}'

    def __init__(self):
        settings = get_project_settings()
        self.mongo_uri = settings.get('MONGODB_URI')
        self.bot_name = settings.get('BOT_NAME')
        self.client = MongoClient(self.mongo_uri)
        self.db_map = dict()
        # 默认初始化当前年的库
        year = str(datetime.datetime.now().year)
        self.db_map[year] = self.client['{}_{}'.format(self.bot_name, year)]

    def get_db(self, batch_no=None, year=None):
        if not batch_no and not year:
            raise Exception('批次编号与年份至少指定一个')
        if not year:
            year = batch_to_year(batch_no)
        db = self.db_map.get(year)
        if not db:
            db = self.client['{}_{}'.format(self.bot_name, year)]
            self.db_map[year] = db

        return db

    def get_batch_data(self, batch_no):
        if batch_no:
            try:
                year = str(batch_no)[:-4]
                db = Mongo().client[self.db_name_tpl.format(year)]
                ids = db['data_{}'.format(batch_no)].find({}, {'_id': 1})
                return list(ids)
            except Exception as ex:
                print('work exception: %s' % format_exc())
                time.sleep(3)
        return []

    def save_orders(self, order_list):
        if order_list:
            try:
                year_month_bulk_map = dict()
                for data in order_list:
                    # 补充采购单位编号
                    order_dept = data.get('orderDept')
                    if order_dept and not data.get('deptId'):
                        data['deptId'] = md5(order_dept)
                    # 计算批次编号
                    batch_no = time_to_batch_no(data.get('orderTime'))
                    if batch_no:
                        # 20190710 -> 201907
                        year_month = str(batch_no)[:-2]
                        to_save = copy.deepcopy(data)
                        bulk_list = year_month_bulk_map.get(year_month, [])
                        bulk_list.append(UpdateOne({'_id': to_save.pop("id")}, {'$set': to_save}, upsert=True))
                        year_month_bulk_map[year_month] = bulk_list
                for year_month, bulk_list in year_month_bulk_map.items():
                    year = str(year_month)[:-2]
                    self.get_db(year=year)['order_item_{}_vvv'.format(year_month)].bulk_write(bulk_list, ordered=False,
                                                                                              bypass_document_validation=True)
            except Exception as ex:
                print('work exception: %s' % format_exc())
                time.sleep(3)
                # 重试
                print('--> retry save_orders: page=%s' % order_list)
                self.save_orders(order_list)
